package com.example.gateway.handler;

import com.alibaba.csp.sentinel.adapter.gateway.sc.callback.BlockRequestHandler;
import com.alibaba.csp.sentinel.adapter.gateway.sc.callback.GatewayCallbackManager;
import com.alibaba.csp.sentinel.slots.block.flow.FlowException;
import com.example.gateway.constant.GatewayExceptionEnum;
import com.example.gateway.exception.DemoEnvForbiddenException;
import com.example.gateway.exception.InvalidTenantException;
import com.example.starter.util.JSONUtil;
import com.example.starter.util.Response;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.boot.web.reactive.error.ErrorWebExceptionHandler;
import org.springframework.cloud.gateway.support.NotFoundException;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.Ordered;
import org.springframework.core.annotation.Order;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.lang.NonNull;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.server.ServerResponse;
import org.springframework.web.server.ResponseStatusException;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

/**
 * <p>
 * 统一异常处理
 * </p>
 *
 * @author 王令
 * @since 2022-05-13 15:14:14
 */
@Configuration
@Order(Ordered.HIGHEST_PRECEDENCE)
@RequiredArgsConstructor
@Slf4j
public class GatewayExceptionHandler implements ErrorWebExceptionHandler, InitializingBean {

    @SneakyThrows
    @Override
    public void afterPropertiesSet() {
        BlockRequestHandler blockRequestHandler = (serverWebExchange, throwable) -> {
            if (throwable instanceof FlowException) {
                return ServerResponse.status(HttpStatus.TOO_MANY_REQUESTS)
                        .contentType(MediaType.APPLICATION_JSON)
                        .body(BodyInserters.fromValue(Response.fail(GatewayExceptionEnum.TOO_MANY_REQUESTS)));
            }
            return ServerResponse
                    .status(HttpStatus.SERVICE_UNAVAILABLE)
                    .contentType(MediaType.APPLICATION_JSON)
                    .body(BodyInserters.empty());
        };
        GatewayCallbackManager.setBlockHandler(blockRequestHandler);
    }

    @Override
    @NonNull
    public Mono<Void> handle(ServerWebExchange exchange, @NonNull Throwable ex) {
        ServerHttpResponse response = exchange.getResponse();
        if (response.isCommitted()) {
            return Mono.error(ex);
        }
        response.getHeaders().setContentType(MediaType.APPLICATION_JSON);
        return response.writeWith(Mono.fromSupplier(() -> {
            DataBufferFactory bufferFactory = response.bufferFactory();
            //设置响应到response的数据
            if (ex instanceof InvalidTenantException) {
                response.setStatusCode(HttpStatus.FORBIDDEN);
                return bufferFactory.wrap(JSONUtil.toJsonByte(Response.fail(GatewayExceptionEnum.INVALID_TENANT)));
            } else if (ex instanceof DemoEnvForbiddenException) {
                response.setStatusCode(HttpStatus.FORBIDDEN);
                return bufferFactory.wrap(JSONUtil.toJsonByte(Response.fail(GatewayExceptionEnum.DEMO_ENV_FORBIDDEN)));
            } else if (ex instanceof NotFoundException) {
                response.setStatusCode(HttpStatus.SERVICE_UNAVAILABLE);
                return bufferFactory.wrap(JSONUtil.toJsonByte(Response.error(GatewayExceptionEnum.SERVICE_UNAVAILABLE)));
            } else if (ex instanceof ResponseStatusException) {
                ResponseStatusException exp = (ResponseStatusException) ex;
                response.setStatusCode(exp.getStatus());
                return bufferFactory.wrap(JSONUtil.toJsonByte(Response.error(exp.getRawStatusCode(), exp.getReason())));
            }
            response.setStatusCode(HttpStatus.INTERNAL_SERVER_ERROR);
            return bufferFactory.wrap(JSONUtil.toJsonByte(Response.error(GatewayExceptionEnum.SYSTEM_ERROR)));
        }));
    }
}
